from typing import Optional

from fastapi import APIRouter, Request, File, UploadFile, Form, Depends, HTTPException, status
from models import ChatMessage, MessageRole
from dependencies import *
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels

from llama_index.indices.managed.dashscope.retriever import DashScopeCloudRetriever

router = APIRouter()

@router.get("/qa")
async def get_qa_page(request: Request):
    return templates.TemplateResponse("qa.html", {"request": request})


@router.post("/qa")
async def post_qa(request: Request,
                  data: dict):

    user_message = data.get("message")
    knowledge_base_name = data.get("knowledge_base_name")
    if not user_message:
        raise HTTPException(status_code=400, detail="请输入问题")
    if not knowledge_base_name:  # 如果没有提供知识库名称，则返回错误
        raise HTTPException(status_code=400, detail="请输入知识库名称")
    # 知识库检索
    retriever = DashScopeCloudRetriever(knowledge_base_name)
    nodes = retriever.retrieve(data["message"])

    # 初始化大模型
    dashscope_llm = DashScope(
        model_name=DashScopeGenerationModels.QWEN_MAX, api_key=DASHSCOPE_API_KEY
    )
    messages = [
        ChatMessage(role=MessageRole.SYSTEM, content="你是一个智能助手，回答问题"),
        ChatMessage(role=MessageRole.USER, content=data["message"])
    ]
    # 添加检索到的信息作为上下文
    for result in nodes:
        text_content = result.node.text
        messages.append(ChatMessage(role=MessageRole.ASSISTANT, content=text_content))
    resp = dashscope_llm.chat(messages)

    print(resp.message.content)
    return {"answer": resp.message.content}


